{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Optical Character Recognition using TAO OCRNet-ViT\n",
    "\n",
    "Transfer learning is the process of transferring learned features from one application to another. It is a commonly used training technique where you use a model trained on one task and re-train to use it on a different task. \n",
    "\n",
    "Train Adapt Optimize (TAO) Toolkit  is a simple and easy-to-use Python based AI toolkit for taking purpose-built AI models and customizing them with users' own data.\n",
    "\n",
    "<img align=\"center\" src=\"https://d29g4g2dyqv443.cloudfront.net/sites/default/files/akamai/TAO/tlt-tao-toolkit-bring-your-own-model-diagram.png\" width=\"1080\">\n",
    "\n",
    "## Sample prediction of OCRNet\n",
    "<img align=\"center\" src=\"https://github.com/vpraveen-nv/model_card_images/blob/main/cv/notebook/ocrnet/OCRNet_inference.png?raw=true\" width=\"960\">"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Learning Objectives\n",
    "In this notebook, you will learn how to leverage the simplicity and convenience of TAO to:\n",
    "\n",
    "* Take a pretrained OCRNet-ViT model and train OCRNet-ViT model on the ICDAR15 dataset\n",
    "* Prune the trained OCRNet-ViT model\n",
    "* Retrain the pruned model to recover lost accuracy\n",
    "* Export the pruned model\n",
    "* Run Inference on the trained model\n",
    "* Export the pruned, and retrained model to a .onnx file for deployment to DeepStream\n",
    "\n",
    "## Table of Contents\n",
    "\n",
    "This notebook shows an example usecase of OCRNet-ViT using Train Adapt Optimize (TAO) Toolkit.\n",
    "\n",
    "0. [Set up env variables and map drives](#head-0)\n",
    "1. [Installing the TAO launcher](#head-1)\n",
    "2. [Prepare dataset and pre-trained model](#head-2) <br>\n",
    "    2.1 [Download pre-trained model](#head-2-1) <br>\n",
    "3. [Provide training specification](#head-3)\n",
    "4. [Run TAO training](#head-4)\n",
    "5. [Evaluate trained models](#head-5)\n",
    "6. [Prune trained models](#head-6)\n",
    "7. [Retrain pruned models](#head-7)\n",
    "8. [Evaluate retrained model](#head-8)\n",
    "9. [Visualize inferences](#head-9)\n",
    "10. [Model Export](#head-10)\n",
    "11. [Verify deployed model](#head-11)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 0. Set up env variables and map drives <a class=\"anchor\" id=\"head-0\"></a>\n",
    "\n",
    "When using the purpose-built pretrained models from NGC, please make sure to set the `$KEY` environment variable to the key as mentioned in the model overview. Failing to do so, can lead to errors when trying to load them as pretrained models.\n",
    "\n",
    "The TAO launcher uses docker containers under the hood, and **for our data and results directory to be visible to the docker, they need to be mapped**. The launcher can be configured using the config file `~/.tao_mounts.json`. Apart from the mounts, you can also configure additional options like the Environment Variables and amount of Shared Memory available to the TAO launcher. <br>\n",
    "\n",
    "`IMPORTANT NOTE:` The code below creates a sample `~/.tao_mounts.json`  file. Here, we can map directories in which we save the data, specs, results and cache. You should configure it for your specific case so these directories are correctly visible to the docker container.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "\n",
    "# Please define this local project directory that needs to be mapped to the TAO docker session.\n",
    "%env LOCAL_PROJECT_DIR=/path/to/local/tao-experiments\n",
    "\n",
    "os.environ[\"HOST_DATA_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"data\", \"ocrnet\")\n",
    "os.environ[\"HOST_RESULTS_DIR\"] = os.path.join(os.getenv(\"LOCAL_PROJECT_DIR\", os.getcwd()), \"ocrnet\")\n",
    "\n",
    "# Set this path if you don't run the notebook from the samples directory.\n",
    "# %env NOTEBOOK_ROOT=/path/to/local/tao-experiments/ocrnet\n",
    "# The sample spec files are present in the same path as the downloaded samples.\n",
    "os.environ[\"HOST_SPECS_DIR\"] = os.path.join(\n",
    "    os.getenv(\"NOTEBOOK_ROOT\", os.getcwd()),\n",
    "    \"specs\"\n",
    ")\n",
    "\n",
    "# Set your encryption key, and use the same key for all commands\n",
    "%env KEY = nvidia_tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! mkdir -p $HOST_DATA_DIR\n",
    "! mkdir -p $HOST_SPECS_DIR\n",
    "! mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Mapping up the local directories to the TAO docker.\n",
    "import json\n",
    "import os\n",
    "mounts_file = os.path.expanduser(\"~/.tao_mounts.json\")\n",
    "tlt_configs = {\n",
    "   \"Mounts\":[\n",
    "       # Mapping the data directory\n",
    "       {\n",
    "           \"source\": os.environ[\"LOCAL_PROJECT_DIR\"],\n",
    "           \"destination\": \"/workspace/tao-experiments\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_DATA_DIR\"],\n",
    "           \"destination\": \"/data\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_SPECS_DIR\"],\n",
    "           \"destination\": \"/specs\"\n",
    "       },\n",
    "       {\n",
    "           \"source\": os.environ[\"HOST_RESULTS_DIR\"],\n",
    "           \"destination\": \"/results\"\n",
    "       },\n",
    "   ],\n",
    "   \"DockerOptions\": {\n",
    "        \"shm_size\": \"16G\",\n",
    "        \"ulimits\": {\n",
    "            \"memlock\": -1,\n",
    "            \"stack\": 67108864\n",
    "         }\n",
    "   }\n",
    "}\n",
    "# Writing the mounts file.\n",
    "with open(mounts_file, \"w\") as mfile:\n",
    "    json.dump(tlt_configs, mfile, indent=4)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat ~/.tao_mounts.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. Installing the TAO launcher <a class=\"anchor\" id=\"head-1\"></a>\n",
    "The TAO launcher is a python package distributed as a python wheel listed in PyPI. You may install the launcher by executing the following cell.\n",
    "\n",
    "Please note that TAO Toolkit recommends users to run the TAO launcher in a virtual env with python 3.6.9. You may follow the instruction in this [page](https://virtualenvwrapper.readthedocs.io/en/latest/install.html) to set up a python virtual env using the `virtualenv` and `virtualenvwrapper` packages. Once you have setup virtualenvwrapper, please set the version of python to be used in the virtual env by using the `VIRTUALENVWRAPPER_PYTHON` variable. You may do so by running\n",
    "\n",
    "```sh\n",
    "export VIRTUALENVWRAPPER_PYTHON=/path/to/bin/python3.x\n",
    "```\n",
    "where x >= 6 and <= 8\n",
    "\n",
    "We recommend performing this step first and then launching the notebook from the virtual environment. In addition to installing TAO python package, please make sure of the following software requirements:\n",
    "* python >=3.7, <=3.10.x\n",
    "* docker-ce > 19.03.5\n",
    "* docker-API 1.40\n",
    "* nvidia-container-toolkit > 1.3.0-1\n",
    "* nvidia-container-runtime > 3.4.0-1\n",
    "* nvidia-docker2 > 2.5.0-1\n",
    "* nvidia-driver > 455+\n",
    "\n",
    "Once you have installed the pre-requisites, please log in to the docker registry nvcr.io by following the command below\n",
    "\n",
    "```sh\n",
    "docker login nvcr.io\n",
    "```\n",
    "\n",
    "You will be triggered to enter a username and password. The username is `$oauthtoken` and the password is the API key generated from `ngc.nvidia.com`. Please follow the instructions in the [NGC setup guide](https://docs.nvidia.com/ngc/ngc-overview/index.html#generating-api-key) to generate your own API key.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SKIP this step IF you have already installed the TAO launcher.\n",
    "!pip3 install nvidia-tao"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# View the versions of the TAO launcher\n",
    "!tao info"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Prepare dataset and pre-trained model <a class=\"anchor\" id=\"head-2\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    " We will be using the ICDAR15 word recognition dataset for the tutorial. To find more details please visit\n",
    "https://rrc.cvc.uab.es/?ch=4&com=tasks. Please download the ICDAR15 word recognition train dataset (https://rrc.cvc.uab.es/?ch=4&com=downloads) to `$HOST_DATA_DIR/train` and test dataset to `$HOST_DATA_DIR/test`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create local dir\n",
    "!mkdir -p $HOST_DATA_DIR\n",
    "!mkdir -p $HOST_RESULTS_DIR"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check the dataset is present\n",
    "!if [ ! -f $HOST_DATA_DIR/test/ch4_test_word_images_gt.zip ]; then echo 'Test Image zip file not found, please download.'; else echo 'Found Test Image zip file.';fi\n",
    "!if [ ! -f $HOST_DATA_DIR/test/Challenge4_Test_Task3_GT.txt ]; then echo 'Test Label file not found, please download.'; else echo 'Found Test Labels file.';fi\n",
    "!if [ ! -f $HOST_DATA_DIR/train/ch4_training_word_images_gt.zip ]; then echo 'Train zip file not found, please download.'; else echo 'Found Train zip file.';fi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# unpack \n",
    "!unzip -u $HOST_DATA_DIR/test/ch4_test_word_images_gt.zip -d $HOST_DATA_DIR/test\n",
    "!unzip -u $HOST_DATA_DIR/train/ch4_training_word_images_gt.zip -d $HOST_DATA_DIR/train"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# verify\n",
    "!ls -l $HOST_DATA_DIR/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The characters_list.txt will contain all the characters found in the dataset. Each character occupies one line. \n",
    "# The following code will process the labels to align with character_list.txt of the pretrained model\n",
    "# Clean the label to alphanumeric, non-sensitive (lower case). Filter the label with length larger than 25\n",
    "import re\n",
    "\n",
    "def preprocess_label(gt_file, filtered_file):\n",
    "    gt_list = open(gt_file, \"r\").readlines()\n",
    "    filtered_list = []\n",
    "\n",
    "    character_list = \"0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ\"\n",
    "    for label_line in gt_list:\n",
    "        try:\n",
    "            path, label = label_line.strip().split()\n",
    "        except Exception:\n",
    "            continue\n",
    "        path = path[:-1]\n",
    "        label = label.strip(\"\\\"\")\n",
    "        if re.search(f\"[^{character_list}]\", label):\n",
    "            continue\n",
    "        else:\n",
    "            if len(label) <= 25:\n",
    "                label = label.lower() # ignore the case\n",
    "                filtered_list.append(f\"{path}\\t{label}\\n\")\n",
    "\n",
    "    with open(filtered_file, \"w\") as f:\n",
    "        f.writelines(filtered_list)\n",
    "\n",
    "orig_train_gt_file=os.path.join(os.getenv(\"HOST_DATA_DIR\"), \"train\", \"gt.txt\")\n",
    "processed_train_gt_file=os.path.join(os.getenv(\"HOST_DATA_DIR\"), \"train\", \"gt_new.txt\")\n",
    "orig_test_gt_file=os.path.join(os.getenv(\"HOST_DATA_DIR\"), \"test\", \"Challenge4_Test_Task3_GT.txt\")\n",
    "processed_test_gt_file=os.path.join(os.getenv(\"HOST_DATA_DIR\"), \"test\", \"gt_new.txt\")\n",
    "preprocess_label(orig_train_gt_file, processed_train_gt_file)\n",
    "preprocess_label(orig_test_gt_file, processed_test_gt_file)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Set the path from the perspective of the TAO docker container\n",
    "%env DATA_DIR = /data\n",
    "%env SPECS_DIR = /specs\n",
    "%env RESULTS_DIR = /results"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Then we will convert the raw dataset (images + labels list) to LMDB format. LMDB is a key-value memory database. With storing the dataset in RAM memory, we can enjoy a better data IO bandwidth. But if we're working with a remote file system which is used by multiple persons at the same time, we should skip the following steps and use raw dataset loader of OCRNet."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the raw train dataset to lmdb\n",
    "print(\"Converting the training set to LMDB.\")\n",
    "!tao model ocrnet dataset_convert -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                            dataset_convert.input_img_dir=$DATA_DIR/train \\\n",
    "                            dataset_convert.gt_file=$DATA_DIR/train/gt_new.txt \\\n",
    "                            dataset_convert.results_dir=$DATA_DIR/train/lmdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert the raw test dataset to lmdb\n",
    "print(\"Converting the testing set to LMDB.\")\n",
    "!tao model ocrnet dataset_convert -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                            dataset_convert.input_img_dir=$DATA_DIR/test \\\n",
    "                            dataset_convert.gt_file=$DATA_DIR/test/gt_new.txt \\\n",
    "                            dataset_convert.results_dir=$DATA_DIR/test/lmdb"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The characters_list.txt will contain all the characters found in the dataset. Each character occupies one line. The model will only classify the characters in this list.\n",
    "# Generate the character list file for the model:\n",
    "character_list = \"!#$%&'()*+,-./0123456789:;<=>[]^_abcdefghijklmnopqrstuvwxyz|~\"\n",
    "with open(os.path.join(os.getenv(\"HOST_DATA_DIR\"), \"character_list\"), \"w\") as f:\n",
    "     for ch in character_list:\n",
    "            f.write(f\"{ch}\\n\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlt $HOST_DATA_DIR/train/lmdb"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Additionally, if you have your own dataset already in a volume (or folder), you can mount the volume on `HOST_DATA_DIR` (or create a soft link). Below shows an example:\n",
    "```bash\n",
    "# if your dataset is in /dev/sdc1\n",
    "mount /dev/sdc1 $HOST_DATA_DIR\n",
    "\n",
    "# if your dataset is in folder /var/dataset\n",
    "ln -sf /var/dataset $HOST_DATA_DIR\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2.1 Download pre-trained model <a class=\"anchor\" id=\"head-2-1\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We will use NGC CLI to get the pre-trained models. For more details, go to [ngc.nvidia.com](ngc.nvidia.com) and click the SETUP on the navigation bar."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Installing NGC CLI on the local machine.\n",
    "## Download and install\n",
    "%env CLI=ngccli_cat_linux.zip\n",
    "!mkdir -p $HOST_RESULTS_DIR/ngccli\n",
    "\n",
    "# Remove any previously existing CLI installations\n",
    "!rm -rf $HOST_RESULTS_DIR/ngccli/*\n",
    "!wget \"https://ngc.nvidia.com/downloads/$CLI\" -P $HOST_RESULTS_DIR/ngccli\n",
    "!unzip -u \"$HOST_RESULTS_DIR/ngccli/$CLI\" -d $HOST_RESULTS_DIR/ngccli/\n",
    "!rm $HOST_RESULTS_DIR/ngccli/*.zip \n",
    "os.environ[\"PATH\"]=\"{}/ngccli/ngc-cli:{}\".format(os.getenv(\"HOST_RESULTS_DIR\", \"\"), os.getenv(\"PATH\", \"\"))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ngc registry model list nvidia/tao/ocrnet:*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/pretrained_ocrnet/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Pull pretrained model from NGC\n",
    "!ngc registry model download-version nvidia/tao/ocrnet:trainable_v2.0 --dest $HOST_RESULTS_DIR/pretrained_ocrnet"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Check that model is downloaded into dir.\")\n",
    "!ls -l $HOST_RESULTS_DIR/pretrained_ocrnet/ocrnet_vtrainable_v2.0"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Provide training specification <a class=\"anchor\" id=\"head-3\"></a>\n",
    "* Dataset for the train datasets\n",
    "    * In order to use the newly generated dataset, update the dataset_config parameter in the spec file at `$HOST_SPECS_DIR/experiment.yaml`\n",
    "    * You also need to prepare the new `charater_list_file`.\n",
    "* Other training (hyper-)parameters such as batch size, number of epochs, learning rate etc."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!cat $HOST_SPECS_DIR/experiment-vit.yaml"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Run TAO training <a class=\"anchor\" id=\"head-4\"></a>\n",
    "* Provide the sample spec file and the output directory location for models\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/experiment_dir_unpruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model ocrnet train -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "              train.results_dir=$RESULTS_DIR/experiment_dir_unpruned \\\n",
    "              train.pretrained_model_path=$RESULTS_DIR/pretrained_ocrnet/ocrnet_vtrainable_v2.0/ocrnet-vit.pth \\\n",
    "              dataset.train_dataset_dir=[$DATA_DIR/train/lmdb] \\\n",
    "              dataset.val_dataset_dir=$DATA_DIR/test/lmdb \\\n",
    "              dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "## Training command for multi-gpu training. We can define the number of gpus and specify which GPU's are to be used by setting the `train.gpu_ids` parameter.\n",
    "## The following command will trigger multi-gpu training on gpu 0 and gpu 1.\n",
    "# !tao model ocrnet train -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "#               train.gpu_ids=[0,1] \\\n",
    "#               train.results_dir=$RESULTS_DIR/experiment_dir_unpruned \\\n",
    "#               train.pretrained_model_path=$RESULTS_DIR/pretrained_ocrnet/ocrnet_vtrainable_v2.0/ocrnet-vit.pth \\\n",
    "#               dataset.train_dataset_dir=[$DATA_DIR/train/lmdb] \\\n",
    "#               dataset.val_dataset_dir=$DATA_DIR/test/lmdb \\\n",
    "#               dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Trained:')\n",
    "print('---------------------')\n",
    "!ls -ltrh $HOST_RESULTS_DIR/experiment_dir_unpruned/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/experiment_dir_unpruned/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"experiment_dir_unpruned/ocr_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/experiment_dir_unpruned/ocr_vit_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/experiment_dir_unpruned/ocr_vit_model.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Evaluate trained models <a class=\"anchor\" id=\"head-5\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model ocrnet evaluate -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                 evaluate.results_dir=$RESULTS_DIR/experiment_dir_unpruned \\\n",
    "                 evaluate.checkpoint=$RESULTS_DIR/experiment_dir_unpruned/ocr_vit_model.pth \\\n",
    "                 evaluate.test_dataset_dir=$DATA_DIR/test/lmdb \\\n",
    "                 dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Prune trained models <a class=\"anchor\" id=\"head-6\"></a>\n",
    "* Specify pre-trained model\n",
    "* Choose the pruning method from [`amount`, `threshold`, `experimental_hybrid`]. Default to be `experimental_hybrid` in this notebook.\n",
    "* `threshold` or `amount` for pruning .\n",
    "* A key to save and load the model\n",
    "* Output directory to store the model\n",
    "\n",
    "Usually, you just need to adjust threshold or amount for accuracy and model size trade off. `amount` is for `amount` and `experimental_hybrid` pruning. The smaller amount, the smaller the pruned model will be. `threshold` is for `threshold` pruning, the higher threshold value, the smaller the pruned model will be. Users can try multiple times to find the best trade-off between the model size and model accruracy. For more details about pruning algorithms, please refer to [TAO-Toolkit documentation]()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/experiment_dir_pruned"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "!tao model ocrnet prune -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "              prune.checkpoint=$RESULTS_DIR/experiment_dir_unpruned/ocr_vit_model.pth \\\n",
    "              prune.results_dir=$RESULTS_DIR/experiment_dir_pruned/ \\\n",
    "              prune.pruned_file=$RESULTS_DIR/experiment_dir_pruned/pruned_0.1.pth \\\n",
    "              dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ls -rlth $HOST_RESULTS_DIR/experiment_dir_pruned/"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Retrain pruned models <a class=\"anchor\" id=\"head-7\"></a>\n",
    "* Model needs to be re-trained to bring back accuracy after pruning\n",
    "* Specify re-training specification\n",
    "* WARNING: training will take several hours or one day to complete"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!mkdir -p $HOST_RESULTS_DIR/experiment_dir_retrain"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Retraining using the pruned model as pretrained weights \n",
    "!tao model ocrnet train -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "              train.results_dir=$RESULTS_DIR/experiment_dir_retrain \\\n",
    "              model.pruned_graph_path=$RESULTS_DIR/experiment_dir_pruned/pruned_0.1.pth \\\n",
    "              dataset.train_dataset_dir=[$DATA_DIR/train/lmdb] \\\n",
    "              dataset.val_dataset_dir=$DATA_DIR/test/lmdb \\\n",
    "              dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Listing the newly retrained model.\n",
    "!ls -rlth $HOST_RESULTS_DIR/experiment_dir_retrain/"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can set NUM_EPOCH to the epoch corresponding to any saved checkpoint\n",
    "# %env NUM_EPOCH=029\n",
    "\n",
    "# Get the name of the checkpoint corresponding to your set epoch\n",
    "# tmp=!ls $HOST_RESULTS_DIR/experiment_dir_retrain/*.pth | grep epoch_$NUM_EPOCH\n",
    "# %env CHECKPOINT={tmp[0]}\n",
    "\n",
    "# Or get the latest checkpoint\n",
    "os.environ[\"CHECKPOINT\"] = os.path.join(os.getenv(\"HOST_RESULTS_DIR\"), \"experiment_dir_retrain/ocr_model_latest.pth\")\n",
    "\n",
    "print('Rename a trained model: ')\n",
    "print('---------------------')\n",
    "!cp $CHECKPOINT $HOST_RESULTS_DIR/experiment_dir_retrain/ocr_vit_model.pth\n",
    "!ls -ltrh $HOST_RESULTS_DIR/experiment_dir_retrain/ocr_vit_model.pth"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Evaluate retrained model <a class=\"anchor\" id=\"head-8\"></a>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tao model ocrnet evaluate -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                 evaluate.results_dir=$RESULTS_DIR/experiment_dir_retrain \\\n",
    "                 evaluate.checkpoint=$RESULTS_DIR/experiment_dir_retrain/ocr_vit_model.pth \\\n",
    "                 evaluate.test_dataset_dir=$DATA_DIR/test/lmdb \\\n",
    "                 model.pruned_graph_path=$RESULTS_DIR/experiment_dir_pruned/pruned_0.1.pth \\\n",
    "                 dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 9. Inferences <a class=\"anchor\" id=\"head-9\"></a>\n",
    "In this section, we run the `infer` tool to generate inferences on the trained models. The predicted label will be printed out in the log."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copy some test images\n",
    "!mkdir -p $HOST_DATA_DIR/test_samples\n",
    "!cp $HOST_DATA_DIR/test/word_100* $HOST_DATA_DIR/test_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Running inference for detection on n images\n",
    "!tao model ocrnet inference -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                  inference.checkpoint=$RESULTS_DIR/experiment_dir_retrain/ocr_vit_model.pth \\\n",
    "                  inference.inference_dataset_dir=$DATA_DIR/test_samples \\\n",
    "                  inference.results_dir=$RESULTS_DIR/experiment_dir_retrain/ \\\n",
    "                  model.pruned_graph_path=$RESULTS_DIR/experiment_dir_pruned/pruned_0.1.pth \\\n",
    "                  dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 10. Model Export <a class=\"anchor\" id=\"head-10\"></a>"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "If you trained a non-QAT model, you may export in FP32, FP16 or INT8 mode using the code block below. For INT8, you need to provide calibration image directory."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "# tao <task> export will fail if .onnx already exists. So we clear the export folder before tao <task> export\n",
    "!rm -rf $HOST_RESULTS_DIR/export\n",
    "# Generate .onnx file using tao container\n",
    "!mkdir -p $HOST_RESULTS_DIR/export\n",
    "\n",
    "# Export the model to .onnx\n",
    "!tao model ocrnet export -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "               export.results_dir=$RESULTS_DIR/export/ \\\n",
    "               export.checkpoint=$RESULTS_DIR/experiment_dir_retrain/ocr_vit_model.pth \\\n",
    "               export.onnx_file=$RESULTS_DIR/export/ocr_vit_model.onnx \\\n",
    "               dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Using the `tao deploy` container, you can generate a TensorRT engine and verify the correctness of the generated through evaluate and inference. \n",
    "\n",
    "The `tao deploy` produces optimized tensorrt engines for the platform that it resides on. Therefore, to get maximum performance, please run `tao deploy` command which will instantiate a deploy container, with the exported `.onnx` file on your target device. The `tao deploy` container only works for x86, with discrete NVIDIA GPU's. \n",
    "\n",
    "For the jetson devices, please download the tao-converter for jetson from the dev zone link [here](https://developer.nvidia.com/tao-converter). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Convert to TensorRT engine (FP32). Change --data_type to fp16 for FP16 mode\n",
    "!tao deploy ocrnet gen_trt_engine -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                               gen_trt_engine.onnx_file=$RESULTS_DIR/export/ocr_vit_model.onnx \\\n",
    "                               gen_trt_engine.trt_engine=$RESULTS_DIR/export/trt.engine \\\n",
    "                               gen_trt_engine.tensorrt.min_batch_size=1 \\\n",
    "                               gen_trt_engine.tensorrt.opt_batch_size=1 \\\n",
    "                               gen_trt_engine.tensorrt.max_batch_size=1 \\\n",
    "                               gen_trt_engine.tensorrt.data_type=fp32"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print('Exported model:')\n",
    "print('------------')\n",
    "!ls -lh $HOST_RESULTS_DIR/export"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 11. Verify the deployed model <a class=\"anchor\" id=\"head-11\"></a>\n",
    "Verify the converted engine by TensorRT inferences."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Infer using TensorRT engine\n",
    "\n",
    "# The engine batch size once created, cannot be alterred. So if you wish to run with a different batch-size,\n",
    "# please re-run tao deploy.\n",
    "\n",
    "!tao deploy ocrnet inference -e $SPECS_DIR/experiment-vit.yaml \\\n",
    "                             inference.trt_engine=$RESULTS_DIR/export/trt.engine \\\n",
    "                             inference.inference_dataset_dir=$DATA_DIR/test_samples \\\n",
    "                             inference.input_width=200 \\\n",
    "                             inference.input_height=64 \\\n",
    "                             dataset.character_list_file=$DATA_DIR/character_list"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Evaluation using TensorRT engine\n",
    "!tao deploy ocrnet evaluate -e $SPECS_DIR/experiment.yaml \\\n",
    "                            evaluate.trt_engine=$RESULTS_DIR/export/trt.engine \\\n",
    "                            evaluate.test_dataset_dir=$DATA_DIR/test \\\n",
    "                            evaluate.test_dataset_gt_file=$DATA_DIR/test/gt_new.txt \\\n",
    "                            evaluate.input_width=200 \\\n",
    "                            evaluate.input_height=64 \\\n",
    "                            dataset.character_list_file=$DATA_DIR/character_list"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
